package scanner

import (
	"context"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"sync"
	"time"

	"ssh_status/internal/v2/core"
	"ssh_status/internal/v2/utils"

	"golang.org/x/crypto/ssh"
)

// Scanner 实现了core.Scanner接口
type Scanner struct {
	config    core.ScannerConfig
	ipNet     *net.IPNet
	ipCount   int
	startTime time.Time
	completed int
	status    string
	sshConfig *ssh.ClientConfig

	// 输出文件
	successFile    *os.File
	sudoFile       *os.File
	portFailedFile *os.File
	authFailedFile *os.File

	// 互斥锁
	fileMutex sync.Mutex
	statMutex sync.Mutex
}

// NewScanner 创建一个新的扫描器实例
func NewScanner(config core.ScannerConfig) (core.Scanner, error) {
	// 验证配置
	if config.Network == "" {
		return nil, fmt.Errorf("网段不能为空")
	}
	if config.Username == "" {
		return nil, fmt.Errorf("用户名不能为空")
	}
	if config.Password == "" {
		return nil, fmt.Errorf("密码不能为空")
	}
	if config.Threads <= 0 {
		config.Threads = 100 // 默认并发数
	}

	// 解析网段
	ipNet, err := utils.ParseIPRange(config.Network)
	if err != nil {
		return nil, err
	}

	// 计算IP数量
	ipCount, err := utils.CalculateIPCount(config.Network)
	if err != nil {
		return nil, err
	}

	// 创建SSH配置
	sshConfig := &ssh.ClientConfig{
		User: config.Username,
		Auth: []ssh.AuthMethod{
			ssh.Password(config.Password),
		},
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         5 * time.Second,
	}

	return &Scanner{
		config:    config,
		ipNet:     ipNet,
		ipCount:   ipCount,
		sshConfig: sshConfig,
		status:    "initialized",
	}, nil
}

// Initialize 实现core.Scanner接口
func (s *Scanner) Initialize(ctx context.Context) error {
	// 创建输出目录
	if err := os.MkdirAll(s.config.OutputDir, 0755); err != nil {
		return fmt.Errorf("创建输出目录失败: %v", err)
	}

	// 创建输出文件
	var err error
	s.successFile, err = os.Create(filepath.Join(s.config.OutputDir, "ssh_success.txt"))
	if err != nil {
		return fmt.Errorf("创建成功文件失败: %v", err)
	}

	s.sudoFile, err = os.Create(filepath.Join(s.config.OutputDir, "ssh_success_sudo.txt"))
	if err != nil {
		return fmt.Errorf("创建sudo文件失败: %v", err)
	}

	s.portFailedFile, err = os.Create(filepath.Join(s.config.OutputDir, "ssh_port_failed.txt"))
	if err != nil {
		return fmt.Errorf("创建端口失败文件失败: %v", err)
	}

	s.authFailedFile, err = os.Create(filepath.Join(s.config.OutputDir, "ssh_auth_failed.txt"))
	if err != nil {
		return fmt.Errorf("创建认证失败文件失败: %v", err)
	}

	return nil
}

// scanIP 扫描单个IP地址
func (s *Scanner) scanIP(ip net.IP) *core.ScanResult {
	result := &core.ScanResult{
		IP:     ip,
		Status: core.StatusUnknown,
	}

	// 检查端口状态
	conn, err := net.DialTimeout("tcp", fmt.Sprintf("%s:22", ip.String()), 5*time.Second)
	if err != nil {
		result.Status = core.StatusPortFailed
		result.Error = err
		return result
	}
	conn.Close()
	result.PortOpen = true
	result.Status = core.StatusPortOpen

	// 尝试SSH连接
	client, err := ssh.Dial("tcp", fmt.Sprintf("%s:22", ip.String()), s.sshConfig)
	if err != nil {
		result.Status = core.StatusAuthFailed
		result.Error = err
		return result
	}
	defer client.Close()

	// 检查root权限
	session, err := client.NewSession()
	if err != nil {
		result.Status = core.StatusSuccess
		return result
	}
	defer session.Close()

	// 尝试执行id命令检查root权限
	err = session.Run("id -u")
	if err == nil {
		result.RootAccess = true
		result.Status = core.StatusRootAccess
	}

	// 检查sudo权限
	err = session.Run("sudo -n true")
	result.Status = core.StatusSuccess
	result.CanSudo = err == nil

	return result
}

// Run 实现core.Scanner接口
func (s *Scanner) Run(ctx context.Context) ([]string, []string, []string, []string, error) {
	s.startTime = time.Now()
	s.status = "running"
	s.completed = 0

	// 创建结果列表
	var success, sudo, portFailed, authFailed []string

	// 创建工作通道
	ipChan := make(chan net.IP, s.config.Threads)
	resultChan := make(chan *core.ScanResult, s.config.Threads)

	// 创建工作组
	var wg sync.WaitGroup

	// 启动工作协程
	for i := 0; i < s.config.Threads; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for ip := range ipChan {
				select {
				case <-ctx.Done():
					return
				default:
					result := s.scanIP(ip)
					resultChan <- result
				}
			}
		}()
	}

	// 启动结果处理协程
	done := make(chan struct{})
	go func() {
		for result := range resultChan {
			s.statMutex.Lock()
			s.completed++
			s.statMutex.Unlock()

			s.fileMutex.Lock()
			switch result.Status {
			case core.StatusSuccess:
				success = append(success, result.IP.String())
				fmt.Fprintf(s.successFile, "%s\n", result.IP.String())
				if result.CanSudo {
					sudo = append(sudo, result.IP.String())
					fmt.Fprintf(s.sudoFile, "%s\n", result.IP.String())
				}
			case core.StatusPortFailed:
				portFailed = append(portFailed, result.IP.String())
				fmt.Fprintf(s.portFailedFile, "%s\n", result.IP.String())
			case core.StatusAuthFailed:
				authFailed = append(authFailed, result.IP.String())
				fmt.Fprintf(s.authFailedFile, "%s\n", result.IP.String())
			case core.StatusPortOpen:
				portFailed = append(portFailed, result.IP.String())
				fmt.Fprintf(s.portFailedFile, "%s\n", result.IP.String())
			case core.StatusRootAccess:
				success = append(success, result.IP.String())
				fmt.Fprintf(s.successFile, "%s\n", result.IP.String())
				if result.RootAccess {
					sudo = append(sudo, result.IP.String())
					fmt.Fprintf(s.sudoFile, "%s\n", result.IP.String())
				}
			}
			s.fileMutex.Unlock()
		}
		close(done)
	}()

	// 生成IP并发送到通道
	ip := s.ipNet.IP
	for i := 0; i < s.ipCount; i++ {
		select {
		case <-ctx.Done():
			close(ipChan)
			<-done
			s.status = "cancelled"
			return success, sudo, portFailed, authFailed, ctx.Err()
		default:
			ipChan <- utils.CopyIP(ip)
			utils.IncrementIP(ip)
		}
	}

	// 关闭IP通道并等待所有工作完成
	close(ipChan)
	wg.Wait()
	close(resultChan)
	<-done

	s.status = "completed"
	return success, sudo, portFailed, authFailed, nil
}

// Status 实现core.Scanner接口
func (s *Scanner) Status() (total int, completed int, status string) {
	s.statMutex.Lock()
	defer s.statMutex.Unlock()
	return s.ipCount, s.completed, s.status
}

// Close 实现core.Scanner接口
func (s *Scanner) Close() error {
	// 关闭所有输出文件
	files := []*os.File{
		s.successFile,
		s.sudoFile,
		s.portFailedFile,
		s.authFailedFile,
	}

	var lastErr error
	for _, f := range files {
		if f != nil {
			if err := f.Close(); err != nil {
				lastErr = err
			}
		}
	}

	return lastErr
}
